import numpy as np
from tqdm import tqdm
import os, sys, yaml, configargparse, pickle, copy, json, h5py, math
from math import cos, sin
import torch
import trimesh
from lib.dif.dif_net import DeformedImplicitField

def cond_mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def rot_X(theta):
    return np.array([[1,0,0],[0,cos(theta),-sin(theta)],[0,sin(theta),cos(theta)]],dtype=np.float32)
def rot_Z(theta):
    return np.array([[cos(theta),-sin(theta),0],[sin(theta),cos(theta),0],[0,0,1]],dtype=np.float32)

def load_grasps(filename):
    """Load transformations and qualities of grasps from a JSON file from the dataset.

    Args:
        filename (str): HDF5 or JSON file name.

    Returns:
        np.ndarray: Homogenous matrices describing the grasp poses. 2000 x 4 x 4.
        np.ndarray: List of binary values indicating grasp success in simulation.
    """
    if filename.endswith(".json"):
        data = json.load(open(filename, "r"))
        T = np.array(data["transforms"])
        success = np.array(data["quality_flex_object_in_gripper"])
    elif filename.endswith(".h5"):
        data = h5py.File(filename, "r")
        T = np.array(data["grasps/transforms"])
        success = np.array(data["grasps/qualities/flex/object_in_gripper"])
    else:
        raise RuntimeError("Unknown file ending:", filename)
    return T, success

def create_gripper_marker(color=[0, 0, 255], tube_radius=0.001, sections=6):
    """Create a 3D mesh visualizing a parallel yaw gripper. It consists of four cylinders.

    Args:
        color (list, optional): RGB values of marker. Defaults to [0, 0, 255].
        tube_radius (float, optional): Radius of cylinders. Defaults to 0.001.
        sections (int, optional): Number of sections of each cylinder. Defaults to 6.

    Returns:
        trimesh.Trimesh: A mesh that represents a simple parallel yaw gripper.
    """
    cfl = trimesh.creation.cylinder(
        radius=0.002,
        sections=sections,
        segment=[
            [4.10000000e-02, -7.27595772e-12, 6.59999996e-02],
            [4.10000000e-02, -7.27595772e-12, 1.12169998e-01],
        ],
    )
    cfr = trimesh.creation.cylinder(
        radius=0.002,
        sections=sections,
        segment=[
            [-4.100000e-02, -7.27595772e-12, 6.59999996e-02],
            [-4.100000e-02, -7.27595772e-12, 1.12169998e-01],
        ],
    )
    cb1 = trimesh.creation.cylinder(
        radius=0.002, sections=sections, segment=[[0, 0, 0], [0, 0, 6.59999996e-02]]
    )
    cb2 = trimesh.creation.cylinder(
        radius=0.002,
        sections=sections,
        segment=[[-4.100000e-02, 0, 6.59999996e-02], [4.100000e-02, 0, 6.59999996e-02]],
    )

    tmp = trimesh.util.concatenate([cb1, cb2, cfr, cfl])
    tmp.visual.face_colors = color

    return tmp

def my_create_gripper_marker(color=[0, 0, 255], tube_radius=0.001, sections=6):
    cfl = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[
            [-7.27595772e-12, -4.10000000e-02, 6.59999996e-02],
            [-7.27595772e-12, -4.10000000e-02, 1.12169998e-01],
        ],
    )
    cfr = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[
            [-7.27595772e-12, 4.10000000e-02, 6.59999996e-02],
            [-7.27595772e-12, 4.10000000e-02, 1.12169998e-01],
        ],
    )
    cb1 = trimesh.creation.cylinder(
        radius=tube_radius, sections=sections, 
        segment=[[0, 0, 0], [0, 0, 6.59999996e-02]],
    )
    cb2 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, -4.100000e-02, 6.59999996e-02], [0, 4.100000e-02, 6.59999996e-02]],
    )

    tmp = trimesh.util.concatenate([cb1, cb2, cfr, cfl])
    tmp.visual.face_colors = color

    return tmp

#############For refine_v2##############
def get_points_sdf_normal(model, shape_code, points):
    sdf = model.inference_with_grad(points, shape_code)
    grad_sdf = torch.autograd.grad(sdf, points, grad_outputs=torch.ones_like(sdf), create_graph=True)[0]
    return sdf.squeeze(0), grad_sdf

def get_template_points_sdf_normal(model, points):
    sdf = model.get_template_field_with_grad(points)
    grad_sdf = torch.autograd.grad(sdf, points, grad_outputs=torch.ones_like(sdf), create_graph=True)[0]
    return sdf.squeeze(0), grad_sdf
########################################

def get_points_normal(model, shape_code, point):
    point = torch.from_numpy(point.astype(np.float32)).cuda()[None, :]
    sdf = model.inference_with_grad(point, shape_code)
    grad_sdf = torch.autograd.grad(sdf, point, grad_outputs=torch.ones_like(sdf), create_graph=True)[0]
    return grad_sdf.squeeze(0).detach().cpu().numpy()

def get_template_points_normal(model, point):
    point = torch.from_numpy(point.astype(np.float32)).cuda()[None, :]
    sdf = model.get_template_field_with_grad(point)
    grad_sdf = torch.autograd.grad(sdf, point, grad_outputs=torch.ones_like(sdf), create_graph=True)[0]
    return grad_sdf.squeeze(0).detach().cpu().numpy()

def vec_norm(vec):
    return vec/np.linalg.norm(vec)

def get_optimize_vector(point_1, point_2, normal_1, normal_2, index):
    vec_1to2 = point_2 - point_1
    vec_1to2_norm = vec_norm(vec_1to2)

    normal_1 = vec_norm(normal_1)
    cos_theta1 = min(1, np.sum(normal_1 * (- vec_1to2_norm)))
    # print(index, " cos theta1: {:.5f}".format(cos_theta1))
    vec_opt1 = vec_1to2_norm + cos_theta1 * normal_1
    # length of point2 to point1 * sin theta1 / 2
    max_step1 = np.linalg.norm(vec_1to2) * ((1 - cos_theta1 ** 2) ** 0.5) / 2

    normal_2 = vec_norm(normal_2)
    cos_theta2 = min(1, np.sum(normal_2 * vec_1to2_norm))
    # print(index, " cos theta2: {:.5f}".format(cos_theta2))
    vec_opt2 = normal_2 * cos_theta2 - vec_1to2_norm
    vec_opt2 = vec_norm(vec_opt2)
    # length of point2 to point1 * sin theta2 / 2
    max_step2 = np.linalg.norm(vec_1to2) * ((1 - cos_theta2 ** 2) ** 0.5) / 2 

    if np.sum(vec_opt1 * vec_opt2) > 0.:
        vec_opt1 = -vec_opt1
        vec_opt2 = -vec_opt2

    if cos_theta1 >= 0.9 and cos_theta2 >= 0.9:
        return None, None, None, None
    elif min(cos_theta1,cos_theta2) < 0.40:
        return vec_opt1, vec_opt2, -1, -1
    else:
        return vec_opt1, vec_opt2, max_step1, max_step2

def optimize_grasp_points(model, left_points, right_points, shape_code=None, step_scale=0.1, vis=False, obj=None, index=0):
    lp = copy.deepcopy(left_points)
    rp = copy.deepcopy(right_points)
    while True: 
        if shape_code is None:
            lp_normal = get_template_points_normal(model, lp)
            rp_normal = get_template_points_normal(model, rp)
        else:
            lp_normal = get_points_normal(model, shape_code, lp)
            rp_normal = get_points_normal(model, shape_code, rp)
        vec_1, vec_2, max_step_1, max_step_2 = get_optimize_vector(lp, rp, lp_normal, rp_normal, index)
    
        # if max_step_1==-1 or index==16:
            # gripper_points_before_refine = trimesh.points.PointCloud([left_points, right_points], colors=[0,0,255])
            # # print(g['left_points'], g['right_points'])
            # gripper_points_after_refine = trimesh.points.PointCloud([lp, rp], colors=[255,0,0])
            # # print(lp, rp)
            # lpn = trimesh.creation.cylinder(radius=0.002, sections=6, segment=[lp, lp+lp_normal*0.1])
            # rpn = trimesh.creation.cylinder(radius=0.002, sections=6, segment=[rp, rp+rp_normal*0.1])
            # lp_opt = trimesh.creation.cylinder(radius=0.002, sections=6, segment=[lp, lp+vec_1*0.1])
            # rp_opt = trimesh.creation.cylinder(radius=0.002, sections=6, segment=[rp, rp+vec_2*0.1])

            # trimesh.Scene([obj, gripper_points_before_refine, gripper_points_after_refine, lpn, rpn, lp_opt, rp_opt]).show()
        if max_step_1==-1:
            return left_points, right_points

        if vec_1 is None:
            break

        lp += vec_1 * max_step_1 * step_scale
        rp += vec_2 * max_step_2 * step_scale

    if vis:
        gripper_points_before_refine = trimesh.points.PointCloud([left_points, right_points], colors=[0,0,255])
        # print(g['left_points'], g['right_points'])
        gripper_points_after_refine = trimesh.points.PointCloud([lp, rp], colors=[255,0,0])
        # print(lp, rp)

        trimesh.Scene([obj, gripper_points_before_refine, gripper_points_after_refine]).show()
    
    return lp, rp

def load_code(model, meta, file_params):
    if file_params['mode']=='train':
        '''get train code'''
        subject_idx_cuda = torch.Tensor([file_params['subject_idx']]).squeeze().long().cuda()[None,...]
        shape_code = model.get_latent_code(subject_idx_cuda)
    elif file_params['mode']=='eval':
        '''get eval code'''
        shape_code = np.loadtxt(os.path.join('dif/eval/{0}'.format(meta['experiment_name']), file_params['filename'], 'checkpoints/embedding_epoch_0049.txt'), dtype=np.float32)
        shape_code = torch.from_numpy(shape_code).cuda()        
    print("latent code shape: ", shape_code.shape)
    return shape_code

def load_recon_ply(meta, file_params):
    if file_params['mode']=='train':
        if file_params['filename'] == 'template':
            obj_recon = trimesh.load('dif/recon/{0}/train/template.ply'.format(meta['experiment_name']))
        else:
            obj_recon = trimesh.load('dif/recon/{0}/train/test{1}.ply'.format(meta['experiment_name'], str(file_params['subject_idx']).zfill(4)))
    elif file_params['mode']=='eval':
        obj_recon = trimesh.load('dif/eval/{0}/{1}/checkpoints/test.ply'.format(meta['experiment_name'], file_params['filename']))
    return obj_recon

def get_model(meta):
    # define DIF-Net
    model = DeformedImplicitField(**meta)
    model.load_state_dict(torch.load(meta['checkpoint_path']))
    return model.cuda()

def R2q(R):
    
    w = R[0,0]+R[1,1]+R[2,2]+1
    x = R[0,0]-R[1,1]-R[2,2]+1
    y = -R[0,0]+R[1,1]-R[2,2]+1
    z = -R[0,0]-R[1,1]+R[2,2]+1
    
    q = np.array([w,x,y,z])

    index = np.argmax(q)

    # assert q[index]>0, "max(q) > 0"
    if q[index]<=0:
        return np.array([0,0,0,1],dtype=np.float32)

    q[index] = math.sqrt(q[index]) / 2

    if index==0:
        q0 = q[index]
        q1 = (R[2,1]-R[1,2]) / (4*q0)
        q2 = (R[0,2]-R[2,0]) / (4*q0)
        q3 = (R[1,0]-R[0,1]) / (4*q0)
    elif index==1:
        q1 = q[index]
        q0 = (R[2,1]-R[1,2]) / (4*q1)
        q2 = (R[0,1]+R[1,0]) / (4*q1)
        q3 = (R[2,0]+R[0,2]) / (4*q1)
    elif index==2:
        q2 = q[index]
        q0 = (R[0,2]-R[2,0]) / (4*q2)
        q1 = (R[0,1]+R[1,0]) / (4*q2)
        q3 = (R[1,2]+R[2,1]) / (4*q2)
    elif index==3:
        q3 = q[index]
        q0 = (R[1,0]-R[0,1]) / (4*q3)
        q1 = (R[2,0]+R[0,2]) / (4*q3)
        q2 = (R[1,2]+R[2,1]) / (4*q3)
    else:
        raise ValueError('index error:' +  str(index))
    # print(np.linalg.norm(np.array([q1,q2,q3,q0],dtype=np.float32)))
    return np.array([q1,q2,q3,q0],dtype=np.float32)

def get_RT_from_grasp_params(panda_grasp_info, obj_scale):
    left_points = panda_grasp_info['left_points']
    right_points = panda_grasp_info['right_points']
    appro = panda_grasp_info['approach_vector']
    depth = panda_grasp_info['depth']

    hori = right_points - left_points
    hori = hori / np.linalg.norm(hori)
    normal = np.cross(hori, appro)
    normal = normal / np.linalg.norm(normal)

    appro_correct = np.cross(normal, hori)

    RT = np.eye(4, dtype=np.float32)
    RT[:3,:3] = np.array([normal, hori, appro_correct], dtype=np.float32).T
    RT[:3,3] = (left_points + right_points) *  (obj_scale / 2) / 2 
    RT[:3,3] = RT[:3,3] - appro_correct * depth

    # hori_correct = np.cross(appro, normal)

    # RT = np.eye(4, dtype=np.float32)
    # RT[:3,:3] = np.array([normal, hori_correct, appro], dtype=np.float32).T
    # RT[:3,3] = (left_points + right_points) * (obj_scale / 2) / 2 
    # RT[:3,3] = RT[:3,3] - appro * depth

    return RT